'''
https://leetcode.cn/problems/find-critical-and-pseudo-critical-edges-in-minimum-spanning-tree
'''
from typing import List


class UF:
    def __init__(self, n):
        self.father = [i for i in range(n)]
        self.sets = n

    def find(self, x):
        if x != self.father[x]:
            self.father[x] = self.find(self.father[x])
        return self.father[x]

    def union(self, x, y):
        fx, fy = self.find(x), self.find(y)
        if fx == fy: return False
        self.father[fy] = fx
        self.sets -= 1
        return True


class Solution:
    # 枚举
    # 关键边：      删除后不能连通 或 删除后总权值和会增加
    # 伪关键边：    可构成最小生成树，且不是唯一的选择(用这个边可以构成最小生成树，不用这个边也可以)
    def findCriticalAndPseudoCriticalEdges(self, n: int, edges: List[List[int]]) -> List[List[int]]:
        uf = UF(n)
        for i, edge in enumerate(edges):
            edge.append(i)      # 记录原下标
        edges.sort(key=lambda x: x[2])

        # 获取最小生成树的权值和
        min_value = 0
        for u, v, w, _ in edges:
            if uf.union(u, v):
                min_value += w
            if uf.sets == 1:
                break
        if uf.sets != 1: return [[],[]]

        m = len(edges)
        critical_edges = []
        pseudo_critical_edges = []
        for i in range(m):
            # 枚举删除 i 这个边，看他是否是关键边
            uf = UF(n)
            value = 0
            for j, (u,v,w, _) in enumerate(edges):
                if i == j: continue
                if uf.union(u, v):
                    value += w
                if uf.sets == 1:
                    break
            if uf.sets != 1 or value > min_value:
                critical_edges.append(edges[i][3])
                continue

            # i 不是关键边，看强行使用它是否能够成最小生成树(看他是否是 pseudo 关键边)
            uf = UF(n)
            uf.union(edges[i][0], edges[i][1])
            value = edges[i][2]
            for j, (u,v,w, _) in enumerate(edges):
                if i == j: continue
                if uf.union(u, v):
                    value += w
                if uf.sets == 1:
                    break
            if uf.sets == 1 and value == min_value:
                pseudo_critical_edges.append(edges[i][3])
        return [critical_edges, pseudo_critical_edges]

# n = 5
# edges = [[0,1,1],[1,2,1],[2,3,2],[0,3,2],[0,4,3],[3,4,3],[1,4,6]]
n = 6
edges = [[0,1,1],[1,2,1],[0,2,1],[2,3,4],[3,4,2],[3,5,2],[4,5,2]]
print(Solution().findCriticalAndPseudoCriticalEdges(n, edges))